import os
import random
from datetime import datetime
from os.path import join

import numpy as np
import torch
from numpy.random import shuffle
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

from ALE.data.dataset_gzsl import Dataset
from ALE.flags import parser
from ALE.models.confg_model import config_model
from common_models import gan
from common_utils.Evaluator import zsl_acc
from common_utils.utils import load_args, l2_normalize

device = 'cuda' if torch.cuda.is_available() else 'cpu'

parser.add_argument('--phase', default='train', help='训练阶段')

"""
    模拟ale 方法的零样本学习
"""
def main():
    print(os.getcwd())
    args = parser.parse_args()
    load_args(args.config, args)
    TIMESTAMP = "{0:%Y-%m-%dT%H-%M-%S/}".format(datetime.now())
    logpath = os.path.join(args.cv_dir, args.name, 'tensorboard/' + args.phase + "/" + TIMESTAMP)
    os.makedirs(logpath, exist_ok=True)
    dataset = Dataset(data_dir=join(args.data_root, args.data_dir), dataset=args.dataset, phase=args.phase)
    val_dataset = Dataset(data_dir=join(args.data_root, args.data_dir), dataset=args.dataset, phase='val')
    model, optimizer = config_model(args=args, dataset=dataset)
    W = torch.rand(dataset.feat_dim, dataset.attr_dim).to(device)
    W = l2_normalize(W.T).T
    beta = np.zeros(len(dataset.train_classes_seen))
    for i in range(1, beta.shape[0]):
        sum_alpha = 0.0
        for j in range(1, i + 1):
            sum_alpha += 1 / j
        beta[i] = sum_alpha

    trainLoader = DataLoader(dataset, batch_size=len(dataset), shuffle=True)
    valLoader = DataLoader(val_dataset, batch_size=len(val_dataset), shuffle=True)
    writer = SummaryWriter(log_dir=logpath, flush_secs=30)
    train_attrs = dataset.train_attrs.to(device)

    # 生成器
    netG = gan.MLP_Generator(input_dim=dataset.attr_dim, output_dim=dataset.feat_dim,
                             layers=[args.ngh, args.ngh * 2, args.ngh]).to(device)
    flag = False
    for epoch in tqdm(range(0, args.epochs + 1), desc="Current epoch"):
        model.train()  # let's witch to training
        # 添加参数position=0 和 leave=True 是进度条不换行
        for idx, data in tqdm(enumerate(trainLoader), total=len(trainLoader), desc=f'Training  epoch {epoch}',
                              position=0, leave=True):
            data = [d.to(device) for d in data]
            imgs, labels, attrs = data
            # 根据图片特征学习对象的属性
            W = update_W(W,imgs=imgs,attrs=attrs,labels=labels,sig = train_attrs,beta= beta,args = args)
            batch_acc = zsl_acc(X = imgs.data.cpu(),W=W.data.cpu(),y_true =labels.data.cpu(),sig = train_attrs.data.cpu())
            print(f"train Accuracy on batch {idx}: {batch_acc}")
        val_loss = 0.0
        for idx, data in tqdm(enumerate(valLoader), total=len(valLoader), desc=f'val  epoch {epoch}', position=0,
                              leave=True):
            data = [d.to(device) for d in data]
            imgs, labels, attrs = data
            batch_acc = zsl_acc(X=imgs.data.cpu(), W=W.data.cpu(), y_true=labels.data.cpu(), sig=dataset.val_attrs.data.cpu())
            print(f"val Accuracy on batch {idx}: {batch_acc}")


def update_W(W, imgs, attrs=None,labels=None, sig=None, beta=None,args = None):
    rand_idx = np.arange(len(imgs))
    shuffle(rand_idx)
    train_classes = torch.unique(labels)
    for j in rand_idx:
        X_n = imgs[j]
        y_n = labels[j]
        y_ = train_classes[train_classes != y_n]
        XW = torch.matmul(X_n, W)
        gt_class_score = torch.matmul(XW, attrs[j])
        for i in range(len(y_)):
            label = random.choice(y_)
            score = 1 + torch.matmul(XW, sig[label]) - gt_class_score
            if score > 0:
                Y = sig[y_n] - sig[label]
                Y = Y.unsqueeze(dim=0)
                temp = X_n.unsqueeze(dim=-1)
                delta = torch.matmul(X_n.unsqueeze(dim=-1), Y)
                be = beta[int(y_.shape[0] / (i + 1))]
                W += args.lr * be * delta
                break
    return W
if __name__ == '__main__':
    main()
